# encoding: utf-8
"""
@author:  liaoxingyu
@contact: sherlockliao01@gmail.com
"""

import torch
from torch import nn

from .backbones.resnet import ResNet, BasicBlock, Bottleneck
from .backbones.resnet_attention import *
from .backbones.senet import SENet, SEResNetBottleneck, SEBottleneck, SEResNeXtBottleneck
from .backbones.resnet_ibn_a import resnet50_ibn_a
from .backbones.mobilev3 import MobileNet_V3_small
from .backbones.resnet_group import ResNet_group
def weights_init_kaiming(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_out')
        nn.init.constant_(m.bias, 0.0)
    elif classname.find('Conv') != -1:
        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')
        if m.bias is not None:
            nn.init.constant_(m.bias, 0.0)
    elif classname.find('BatchNorm') != -1:
        if m.affine:
            nn.init.constant_(m.weight, 1.0)
            nn.init.constant_(m.bias, 0.0)


def weights_init_classifier(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.normal_(m.weight, std=0.001)
        if m.bias:
            nn.init.constant_(m.bias, 0.0)


class Baseline(nn.Module):
    in_planes = 2048

    def __init__(self, num_classes, last_stride, model_path, neck, neck_feat, model_name, pretrain_choice):
        super(Baseline, self).__init__()
        if model_name == 'resnet18':
            self.in_planes = 512
            self.base = ResNet(last_stride=last_stride, 
                               block=BasicBlock, 
                               layers=[2, 2, 2, 2])
        elif model_name == 'resnet18_os':
            self.in_planes = 512
            from .backbones.resnet_inception import ResNet1
            self.base = ResNet1(last_stride=last_stride,
                               block=BasicBlock,
                               layers=[2, 2, 2, 2])
        elif model_name == 'googlenet':
            self.in_planes = 1024
            import torchvision
            googlenet = torchvision.models.googlenet(pretrained=True)
            self.base = nn.Sequential(*list(googlenet.children())[:-2])
        elif model_name == 'resnet18_group':
            self.in_planes = 512
            self.base = ResNet_group(last_stride=last_stride,
                               layers=[2, 2, 2, 2])
        elif model_name == 'resnet18_group_attention':
            self.in_planes = 512
            self.base = ResNet_CBAM_group(last_stride=last_stride,
                               layers=[2, 2, 2, 2])
        # elif model_name == 'resnet18_cbam':
        #     self.in_planes = 512
        #     self.base = ResNet_CBAM(last_stride=last_stride,
        #                        block=BasicBlock,
        #                        layers=[2, 2, 2, 2])
        elif model_name == 'osnet_x1_0':
            from .backbones.osnet import osnet_x1_0
            self.in_planes = 512
            self.base = osnet_x1_0()
        elif model_name == 'osnetgc_x1_0':
            from .backbones.osnet_gc import osnetgc_x1_0
            self.in_planes = 512
            self.base = osnetgc_x1_0()
        elif model_name == 'osnetgc[2(02)00]_x1_0':
            from .backbones.osnet_gc2 import osnetgc_x1_0
            self.in_planes = 512
            self.base = osnetgc_x1_0()
        elif model_name == 'osnetgcfby_x1_0':
            from .backbones.osnet_gc_fby import osnetgc_x1_0
            self.in_planes = 512
            self.base = osnetgc_x1_0()
        elif model_name == 'osnetgcfby2_x1_0':
            from .backbones.osnet_gc_fby2 import osnet_x1_0
            self.in_planes = 512
            self.base = osnet_x1_0()
        elif model_name == 'osnetgcfby2_1_x1_0':
            from .backbones.osnet_gc_fby2_1 import osnet_x1_0
            self.in_planes = 512
            self.base = osnet_x1_0()
        elif model_name == 'osnet_x0_75':
            from .backbones.osnet import osnet_x0_75
            self.in_planes = 384
            self.base = osnet_x0_75()
        elif model_name == 'osnet_x0_5':
            from .backbones.osnet import osnet_x0_5
            self.in_planes = 256
            self.base = osnet_x0_5()
        elif model_name == 'osnet_x0_25':
            from .backbones.osnet import osnet_x0_25
            self.in_planes = 128
            self.base = osnet_x0_25()
        elif model_name =="inceptionv4":
            from .backbones.inceptionv4 import Inceptionv4
            self.in_planes = 1536
            self.base =Inceptionv4()
        elif model_name =="HACNN":
            from .backbones.HACNN import HACNN
            self.in_planes = 384
            self.base = HACNN()
        elif model_name == 'resnet50_nl':
            from .backbones.resnet_nl import ResNetNL
            self.base = ResNetNL(last_stride=last_stride,
                               block=Bottleneck,
                               layers=[3, 4, 6, 3],
                               non_layers=[0, 2, 3, 0])
        elif model_name == 'resnet18_nl':
            self.in_planes = 512
            from .backbones.resnet_nl import ResNetNL
            self.base = ResNetNL(last_stride=last_stride,
                               block=BasicBlock,
                               layers=[2, 2, 2, 2],
                               non_layers=[0, 2, 2, 0])
        elif model_name == 'resnet18_gc':
            self.in_planes = 512
            from .backbones.resnet_gc import ResNetGC
            self.base = ResNetGC(last_stride=last_stride,
                               block=BasicBlock,
                               layers=[2, 2, 2, 2],
                               non_layers=[0, 2, 2, 0])
        elif model_name == 'resnet50_gc':
            from .backbones.resnet_gc import ResNetGC
            self.base = ResNetGC(last_stride=last_stride,
                               block=Bottleneck,
                               layers=[3, 4, 6, 3],
                               non_layers=[0, 2, 3, 0])
        elif model_name == 'shufflenet':
            import torchvision
            aa = torchvision.models.shufflenet_v2_x1_0(pretrained=True)
            self.base = nn.Sequential(*list(aa.children())[:-1])
            self.in_planes = 1024
        elif model_name == 'Xception':
            self.in_planes = 728
            from .backbones.Xception import Xception
            self.base = Xception()
        elif model_name == 'resnet34':
            self.in_planes = 512
            self.base = ResNet(last_stride=last_stride,
                               block=BasicBlock,
                               layers=[3, 4, 6, 3])
        elif model_name == 'resnet50':
            self.base = ResNet(last_stride=last_stride,
                               block=Bottleneck,
                               layers=[3, 4, 6, 3])
        elif model_name == 'resnet101':
            self.base = ResNet(last_stride=last_stride,
                               block=Bottleneck, 
                               layers=[3, 4, 23, 3])
        elif model_name == 'resnet152':
            self.base = ResNet(last_stride=last_stride, 
                               block=Bottleneck,
                               layers=[3, 8, 36, 3])
            
        elif model_name == 'se_resnet50':
            self.base = SENet(block=SEResNetBottleneck, 
                              layers=[3, 4, 6, 3], 
                              groups=1, 
                              reduction=16,
                              dropout_p=None, 
                              inplanes=64, 
                              input_3x3=False,
                              downsample_kernel_size=1, 
                              downsample_padding=0,
                              last_stride=last_stride) 
        elif model_name == 'se_resnet101':
            self.base = SENet(block=SEResNetBottleneck, 
                              layers=[3, 4, 23, 3], 
                              groups=1, 
                              reduction=16,
                              dropout_p=None, 
                              inplanes=64, 
                              input_3x3=False,
                              downsample_kernel_size=1, 
                              downsample_padding=0,
                              last_stride=last_stride)
        elif model_name == 'se_resnet152':
            self.base = SENet(block=SEResNetBottleneck, 
                              layers=[3, 8, 36, 3],
                              groups=1, 
                              reduction=16,
                              dropout_p=None, 
                              inplanes=64, 
                              input_3x3=False,
                              downsample_kernel_size=1, 
                              downsample_padding=0,
                              last_stride=last_stride)  
        elif model_name == 'se_resnext50':
            self.base = SENet(block=SEResNeXtBottleneck,
                              layers=[3, 4, 6, 3], 
                              groups=32, 
                              reduction=16,
                              dropout_p=None, 
                              inplanes=64, 
                              input_3x3=False,
                              downsample_kernel_size=1, 
                              downsample_padding=0,
                              last_stride=last_stride) 
        elif model_name == 'se_resnext101':
            self.base = SENet(block=SEResNeXtBottleneck,
                              layers=[3, 4, 23, 3], 
                              groups=32, 
                              reduction=16,
                              dropout_p=None, 
                              inplanes=64, 
                              input_3x3=False,
                              downsample_kernel_size=1, 
                              downsample_padding=0,
                              last_stride=last_stride)
        elif model_name == 'senet154':
            self.base = SENet(block=SEBottleneck, 
                              layers=[3, 8, 36, 3],
                              groups=64, 
                              reduction=16,
                              dropout_p=0.2, 
                              last_stride=last_stride)
        elif model_name == 'resnet50_ibn_a':
            self.base = resnet50_ibn_a(last_stride)
        elif model_name == 'MobileNetv3_small':
            self.in_planes = 1024
            self.base = MobileNet_V3_small()
            pass
        # elif model_name == 'MobileNetv3_large':
        #     self.in_planes = 1024
        #     self.base = MobileNet_V3_large()
        if not ((model_name == 'googlenet') or (model_name == 'osnet_x1_0')
                or (model_name == 'osnet_x0_75') or (model_name == 'osnet_x0_25') or (model_name == 'osnet_x0_5')
                or(model_name == 'osnetgc_x1_0') or (model_name == 'osnetgcfby_x1_0')or (model_name == 'osnetgcfby2_x1_0')
            or ("osnet" in model_name) or ("shuffle" in model_name)
        ):
            if pretrain_choice == 'imagenet':
                print('Loading pretrained ImageNet model......')
                self.base.load_param(model_path)

            elif pretrain_choice == 'random':
                print('random init model......')
                self.base.random_init()

        # self.gap = nn.AdaptiveAvgPool2d(1)
        from layers.gem_pool import GeneralizedMeanPoolingP
        self.gap = GeneralizedMeanPoolingP()
        self.num_classes = num_classes
        self.neck = neck
        self.neck_feat = neck_feat

        if self.neck == 'no':
            self.classifier = nn.Linear(self.in_planes, self.num_classes)
            # self.classifier = nn.Linear(self.in_planes, self.num_classes, bias=False)     # new add by luo
            # self.classifier.apply(weights_init_classifier)  # new add by luo
        elif self.neck == 'bnneck':
            self.bottleneck = nn.BatchNorm1d(self.in_planes)
            self.bottleneck.bias.requires_grad_(False)  # no shift
            self.classifier = nn.Linear(self.in_planes, self.num_classes, bias=False)

            self.bottleneck.apply(weights_init_kaiming)
            self.classifier.apply(weights_init_classifier)
        # _______________________________________________________________--
    # def forward(self, x):
    #
    #     global_feat = self.gap(self.base(x))  # (b, 2048, 1, 1)
    #     global_feat = global_feat.view(global_feat.shape[0], -1)  # flatten to (bs, 2048)
    #
    #     if self.neck == 'no':
    #         feat = global_feat
    #     elif self.neck == 'bnneck':
    #         feat = self.bottleneck(global_feat)  # normalize for angular softmax
    #
    #     # cls_score = self.classifier(feat)
    #     # return cls_score
    #
    #     if self.training:
    #         cls_score = self.classifier(feat)
    #         return cls_score, global_feat  # global feature for triplet loss
    #     else:
    #         if self.neck_feat == 'after':
    #             # print("Test with feature after BN")
    #             return feat
    #         else:
    #             # print("Test with feature before BN")
    #             return global_feat

        # _______________________________________________________________--

        self.branch1 = MultiScaleLayer_v2(128, num_classes)
        self.branch2 = MultiScaleLayer_v2(256, num_classes)

    def forward(self, x):

        ###
        x = self.base.conv1(x)
        x = self.base.bn1(x)
        x = self.base.maxpool(x)

        x = self.base.layer1(x)
        x = self.base.attention1(x)
        x = self.base.layer2(x)
        if self.training:
            cls_score_br1 = self.branch1(x)
        x = self.base.layer3(x)
        if self.training:
            cls_score_br2 = self.branch2(x)
        x = self.base.layer4(x)
        # self.base(x)

        global_feat = self.gap(x)  # (b, 2048, 1, 1)
        global_feat = global_feat.view(global_feat.shape[0], -1)  # flatten to (bs, 2048)

        if self.neck == 'no':
            feat = global_feat
        elif self.neck == 'bnneck':
            feat = self.bottleneck(global_feat)  # normalize for angular softmax

        # cls_score = self.classifier(feat)
        # return cls_score

        if self.training:
            cls_score = self.classifier(feat)
            lf_xent = []
            lf_xent.extend([cls_score, cls_score_br1, cls_score_br2])
            return tuple(lf_xent), global_feat  # global feature for triplet loss
        else:
            if self.neck_feat == 'after':
                # print("Test with feature after BN")
                return feat
            else:
                # print("Test with feature before BN")
                return global_feat
            #####+++++++++++++++++++++++++++++++++++++++++
    def load_param(self, trained_path):
        # param_dict = torch.load(trained_path)
        param_dict = torch.load(trained_path).state_dict()
        for i in param_dict:
            if 'classifier' in i:
                continue
            self.state_dict()[i].copy_(param_dict[i])
